import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
from src.util.trajectory_manager import TrajectoryManager
from src.util.save_mp4 import save_mp4
import os   # 仅保留一次


class StarAnimation:
    """
    通用多体轨迹动画
    参数含义与原始版本完全一致，此处省略重复注释
    """
    def __init__(self, *, bodies_info, system, dt=0.1, frames=300,
                 interval=50, fps=20, xlim=(-5, 5), ylim=(-5, 5),
                 title='Star System Simulation',
                 mp4_path='out/simulation.mp4',
                 max_trace=100, max_trail=30):
        self.bodies_info = bodies_info
        self.system      = system
        self.dt          = dt
        self.frames      = frames
        self.interval    = interval
        self.fps         = fps
        self.xlim        = xlim
        self.ylim        = ylim
        self.title       = title
        self.mp4_path    = mp4_path

        # 轨迹管理器
        self.traj_mgr = TrajectoryManager(max_trace, max_trail)
        for b in bodies_info:
            self.traj_mgr.add_body(b['name'])

        # 画布样式
        self.fig, self.ax = plt.subplots(figsize=(8, 8))
        self.fig.set_facecolor('black')
        self.ax.set_facecolor('black')
        self.ax.set_xlim(*xlim)
        self.ax.set_ylim(*ylim)
        self.ax.set_aspect('equal')
        self.ax.set_title(title, color='white')
        self.ax.tick_params(colors='white')
        for spine in self.ax.spines.values():
            spine.set_color('white')

        # 绘图对象
        self.plots = {}
        for b in bodies_info:
            color = b['color']
            name  = b['name']
            ms    = b['radius']
            self.plots[name] = {
                'body':  self.ax.plot([], [], 'o', color=color, markersize=ms, label=name)[0],
                'trace': self.ax.plot([], [], '-', color=color, alpha=0.7, lw=1)[0],
                'trail': self.ax.plot([], [], '-', color=color, alpha=0.3, lw=0.5)[0],
            }
        self.ax.legend()

        # 快速映射：name -> 实体
        self._map = {}
        if hasattr(system, 'star'):   # 单星
            self._map[system.star.name] = system.star
        elif hasattr(system, 'stars'):  # 双星/多星
            for s in system.stars:
                self._map[s.name] = s

    def _animate(self, idx):
        # 1. 物理积分
        self.system.evolve(self.dt)

        # 2. 更新轨迹
        for b in self.bodies_info:
            name = b['name']
            body = self._map[name]
            self.traj_mgr.update_position(name, body.position[0], body.position[1])

        # 3. 刷新图形
        artists = []
        for name, p in self.plots.items():
            body = self._map[name]
            x, y = body.position[0], body.position[1]

            p['body'].set_data([x], [y])
            tx, ty = self.traj_mgr.get_trace_data(name)
            p['trace'].set_data(tx, ty)
            trx, try_ = self.traj_mgr.get_trail_data(name)
            if len(trx) > 1:
                p['trail'].set_data(trx, try_)
            artists.extend([p['body'], p['trace'], p['trail']])
        return artists

    def run(self, show=True):
        """
        运行动画生成和保存流程
        :param show (bool) 是否显示动画图形界面，默认为True
        """
        # 创建动画对象
        anim = FuncAnimation(self.fig, self._animate,
                             frames=self.frames,
                             interval=self.interval,
                             blit=True)

        # 确保保存路径的目录存在
        os.makedirs(os.path.dirname(self.mp4_path), exist_ok=True)

        # 保存动画为MP4文件
        save_mp4(anim, self.mp4_path, fps=self.fps)

        # 根据参数决定是否显示图形界面
        if show:
            plt.show()
